#include "pefileloader.h"
#include <QFile>
#include <QListIterator>
#include <QTextStream>


PEFile::PEFile()
{
}

char *PEFile::getHandle()
{
    return this->pMemoryMappedImage;
}

uint PEFile::getSize()
{
    return this->m_pPEHeader->OptionalHeader.SizeOfImage;
}

FileSection *PEFile::getSection(QString name)
{
    QListIterator<FileSection*> sectionIter(this->m_sections);
    while(sectionIter.hasNext())
    {
        FileSection *section = sectionIter.next();
        if(section->getSectionName() == name)
        {
            return section;
        }
    }
    return NULL;
}

QList<FileSection*> PEFile::getSections()
{
    return this->m_sections;
}

//load into memory like a loader would to it
bool PEFile::load(QString fileName)
{
    QFile file(fileName);
    if(!file.open(QIODevice::ReadOnly))
    {
        //Could not open file
        return false;
    }

    //only pointer to PE header structure important, rest just stub for DOS executable
    IMAGE_DOS_HEADER *dosHeader = new IMAGE_DOS_HEADER();
    file.read((char*)dosHeader, sizeof(IMAGE_DOS_HEADER));
    if(dosHeader->e_magic != IMAGE_DOS_SIGNATURE)
    {
        //not valid PE-file
    }

    //jump to PE header structure
    file.seek(dosHeader->e_lfanew);

    //read signature, image file header and optional header
    IMAGE_NT_HEADERS32 *peHeader = new IMAGE_NT_HEADERS32();
    file.read((char*)peHeader, sizeof(IMAGE_NT_HEADERS32));
    if(peHeader->Signature != IMAGE_NT_SIGNATURE)
    {
        //not valid PE-file
    }

    //preferred address at which we want to load exe, if cant, we need to modify relocation stuff
    uint baseAddressImage = peHeader->OptionalHeader.ImageBase;
    //allocate space in memory to map sections to it like a loader
    uint imageSize = peHeader->OptionalHeader.SizeOfImage;
    pMemoryMappedImage = NULL;
    try
    {
        pMemoryMappedImage = new char[imageSize];
    }
    catch(std::bad_alloc& ba)
    {
        ba.what();
        return false;
    }

    //copy the dos header to start of memory image
    m_pDosHeader = reinterpret_cast<IMAGE_DOS_HEADER*>(memcpy(pMemoryMappedImage, reinterpret_cast<void*>(dosHeader), sizeof(IMAGE_DOS_HEADER)));

    //copy pe header
    m_pPEHeader = reinterpret_cast<IMAGE_NT_HEADERS32*>(memcpy(pMemoryMappedImage + dosHeader->e_lfanew, reinterpret_cast<void*>(peHeader), sizeof(IMAGE_NT_HEADERS32)));
    delete dosHeader;
    delete peHeader;

    //read and copy section tables
    ushort numSections = m_pPEHeader->FileHeader.NumberOfSections;
    file.read((char*)(m_pPEHeader + sizeof(IMAGE_NT_HEADERS32)), numSections * sizeof(IMAGE_SECTION_HEADER));
    IMAGE_SECTION_HEADER *sectionHeaders = (IMAGE_SECTION_HEADER *)(m_pPEHeader + sizeof(IMAGE_NT_HEADERS32));

    //copy sections to right memory image locations
    for(ushort i=0; i < numSections; i++)
    {
        IMAGE_SECTION_HEADER sectionHeader = sectionHeaders[i];
        uint virtualStartAddress = sectionHeader.VirtualAddress;
        uint size = sectionHeader.Misc.VirtualSize;
        //with non-rounded (towards alignment) size of section
        uint virtualEndAddress = virtualStartAddress + size - 1;
        //The section start address inside the exe
        uint rawStartAddress = sectionHeader.PointerToRawData;
        //copy the section to the virtual address of image
        file.seek(sectionHeader.PointerToRawData);
        file.read(pMemoryMappedImage + sectionHeader.VirtualAddress, size);

        PESection *peSection = new PESection(&sectionHeaders[i], pMemoryMappedImage + sectionHeader.VirtualAddress);
        m_sections.append(peSection);

        if(peSection->getSectionName() == ".idata")
        {
            //read import table
            file.seek(rawStartAddress);
            uint dllNameOffset;
            do
            {
                IMAGE_IMPORT_DESCRIPTOR *importDescriptor = new IMAGE_IMPORT_DESCRIPTOR();
                file.read((char*)importDescriptor, sizeof(IMAGE_IMPORT_DESCRIPTOR));
                //pointer to ascii dll name
                dllNameOffset = importDescriptor->Name;
            } while(dllNameOffset != NULL);
        }
        else if(peSection->getSectionName() == ".idata")
        {
            //read export table
            file.seek(rawStartAddress);
            IMAGE_EXPORT_DIRECTORY *exportDescriptor = new IMAGE_EXPORT_DIRECTORY();
            file.read((char*)exportDescriptor, sizeof(IMAGE_EXPORT_DIRECTORY));
            uint numExportedFunctions = exportDescriptor->NumberOfFunctions;
            uint functionAddressTable = exportDescriptor->AddressOfFunctions;
        }
    }
    emit fileLoaded(this);
    return true;
}
